{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "\n",
    "\n",
    "from __future__ import division\n",
    "\n",
    "\n",
    "import os, sys, shutil, time, random\n",
    "import argparse\n",
    "import torch\n",
    "import torch.nn as nn\n",
    "import torch.nn.functional as F\n",
    "import torch.backends.cudnn as cudnn\n",
    "import torchvision.datasets as dset \n",
    "import torchvision.transforms as transforms\n",
    "import time\n",
    "import moxing as mox\n",
    "import math\n",
    "import numpy as np\n",
    "import pickle\n",
    "from scipy.spatial import distance\n",
    "import pdb\n",
    "import copy\n",
    "from torch.nn.parameter import Parameter\n",
    "import warnings\n",
    "warnings.filterwarnings(\"ignore\")\n",
    "mox.file.shift('os','mox')\n",
    "\n",
    "os.chdir('./SCOP_NeurIPS2020/')\n",
    "from utils import AverageMeter, RecorderMeter, time_string, convert_secs2time, timing\n",
    "import models\n",
    "from models import resnet_imagenet, generator_cifar,generator_imagenet\n",
    "from pruning_modules import Kf_Conv2d,Masked_Conv2d_bn,Pruned_Conv2d_bn1,Pruned_Conv2d_bn_middle,Pruned_Conv2d_bn2\n",
    "\n",
    "\n",
    "parser = argparse.ArgumentParser(description='SCOP',\n",
    "                                 formatter_class=argparse.ArgumentDefaultsHelpFormatter)\n",
    "parser.add_argument('--data_path', type=str, default='/cache/tyh/cifar10/',help='Path to dataset')\n",
    "parser.add_argument('--dataset', type=str,default='imagenet', choices=['cifar10','imagenet'])\n",
    "parser.add_argument('--arch',  default='resnet50')\n",
    "# Optimization options\n",
    "\n",
    "parser.add_argument('--batch_size', type=int, default=1024, help='Batch size.')\n",
    "parser.add_argument('--batch_size_kf', type=int, default=512, help='Batch size.')\n",
    "\n",
    "parser.add_argument('--epochs', type=int, default=20, help='Number of epochs to train.')\n",
    "parser.add_argument('--learning_rate', type=float, default=0.004, help='The Learning Rate.')\n",
    "\n",
    "parser.add_argument('--momentum', type=float, default=0.9, help='Momentum.')\n",
    "parser.add_argument('--decay', type=float, default=0.0001, help='Weight decay (L2 penalty).')\n",
    "\n",
    "parser.add_argument('--print_freq', default=200, type=int, metavar='N', help='print frequency (default: 200)')\n",
    "parser.add_argument('--save_path', type=str, default='./tmp/', help='Folder to save checkpoints and log.')\n",
    "\n",
    "parser.add_argument('--evaluate',type=int,default=1, help='evaluate model on validation set')\n",
    "\n",
    "parser.add_argument('--ngpu', type=int, default=1)\n",
    "parser.add_argument('--workers', type=int, default=2, help='number of data loading workers (default: 2)')\n",
    "\n",
    "parser.add_argument('--manualSeed', type=int, help='manual seed')\n",
    "\n",
    "\n",
    "parser.add_argument('--prune_rate', type=float, default=0.4, help='the reducing ratio of pruning based on knockoff')\n",
    "\n",
    "\n",
    "parser.add_argument('--epochs_ft', type=int, default=120)\n",
    "parser.add_argument('--lr_ft', type=float, default=0.2, help='The Learning Rate.')\n",
    "\n",
    "parser.add_argument('--pretrain_path', default='', type=str, help='..path of pre-trained model')\n",
    "\n",
    "\n",
    "\n",
    "\n",
    "\n",
    "args,unparsed = parser.parse_known_args()\n",
    "\n",
    "\n",
    "\n",
    "args.use_cuda = args.ngpu > 0 and torch.cuda.is_available()\n",
    "\n",
    "if args.manualSeed is None:\n",
    "    args.manualSeed = random.randint(1, 10000)\n",
    "random.seed(args.manualSeed)\n",
    "torch.manual_seed(args.manualSeed)\n",
    "if args.use_cuda:\n",
    "    torch.cuda.manual_seed_all(args.manualSeed)\n",
    "cudnn.benchmark = True\n",
    "\n",
    "\n",
    "if args.dataset=='imagenet':\n",
    "    args.workers=16"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def train(train_loader, model, criterion, optimizer, epoch, log):\n",
    "    batch_time = AverageMeter()\n",
    "    data_time = AverageMeter()\n",
    "    losses = AverageMeter()\n",
    "    top1 = AverageMeter()\n",
    "    top5 = AverageMeter()\n",
    "    # switch to train mode\n",
    "    model.train()\n",
    "\n",
    "    end = time.time()\n",
    "    for i, (input, target) in enumerate(train_loader):\n",
    "        # measure data loading time\n",
    "        data_time.update(time.time() - end)\n",
    "\n",
    "        if args.use_cuda:\n",
    "            target = target.cuda(async=True)\n",
    "            input = input.cuda()\n",
    "        input_var = torch.autograd.Variable(input)\n",
    "        target_var = torch.autograd.Variable(target)\n",
    "\n",
    "        output = model(input_var)\n",
    "        loss = criterion(output, target_var)\n",
    "\n",
    "\n",
    "        prec1, prec5 = accuracy(output.data, target, topk=(1, 5))\n",
    "        losses.update(loss.data.item(), input.size(0))\n",
    "        top1.update(prec1.item(), input.size(0))\n",
    "        top5.update(prec5.item(), input.size(0))\n",
    "\n",
    "    \n",
    "        optimizer.zero_grad()\n",
    "        loss.backward()\n",
    "        optimizer.step()\n",
    "\n",
    "\n",
    "        batch_time.update(time.time() - end)\n",
    "        end = time.time()\n",
    "\n",
    "        if i % args.print_freq == 0:\n",
    "            print('  Epoch: [{:03d}][{:03d}/{:03d}]   '\n",
    "                      'Time {batch_time.val:.3f} ({batch_time.avg:.3f})   '\n",
    "                      'Data {data_time.val:.3f} ({data_time.avg:.3f})   '\n",
    "                      'Loss {loss.val:.4f} ({loss.avg:.4f})   '\n",
    "                      'Prec@1 {top1.val:.3f} ({top1.avg:.3f})   '\n",
    "                      'Prec@5 {top5.val:.3f} ({top5.avg:.3f})   '.format(\n",
    "                epoch, i, len(train_loader), batch_time=batch_time,\n",
    "                data_time=data_time, loss=losses, top1=top1, top5=top5) + time_string())\n",
    "    print(\n",
    "        '  **Train** Prec@1 {top1.avg:.3f} Prec@5 {top5.avg:.3f} Error@1 {error1:.3f}'.format(top1=top1, top5=top5,\n",
    "                                                                                              error1=100 - top1.avg))\n",
    "    return top1.avg, losses.avg\n",
    "\n",
    "\n",
    "def validate(val_loader, model, criterion, log):\n",
    "    losses = AverageMeter()\n",
    "    top1 = AverageMeter()\n",
    "    top5 = AverageMeter()\n",
    "\n",
    "\n",
    "    model.eval()\n",
    "\n",
    "    for i, (input, target) in enumerate(val_loader):\n",
    "        if args.use_cuda:\n",
    "            target = target.cuda(async=True)\n",
    "            input = input.cuda()\n",
    "        input_var = torch.autograd.Variable(input)\n",
    "        target_var = torch.autograd.Variable(target)\n",
    "\n",
    "        # compute output\n",
    "        with torch.no_grad():\n",
    "            output = model(input_var)\n",
    "            loss = criterion(output, target_var)\n",
    "\n",
    "        # measure accuracy and record loss\n",
    "        prec1, prec5 = accuracy(output.data, target, topk=(1, 5))\n",
    "        losses.update(loss.data.item(), input.size(0))\n",
    "        top1.update(prec1.item(), input.size(0))\n",
    "        top5.update(prec5.item(), input.size(0))\n",
    "\n",
    "    print('  **Test** Prec@1 {top1.avg:.3f} Prec@5 {top5.avg:.3f} Error@1 {error1:.3f}'.format(top1=top1, \n",
    "                                                                                                   top5=top5,error1=100 - top1.avg))\n",
    "\n",
    "    return top1.avg, top5.avg, losses.avg\n",
    "\n",
    "\n",
    "def save_checkpoint(state, is_best, save_path, filename):\n",
    "    filename = os.path.join(save_path, filename)\n",
    "    torch.save(state, filename)\n",
    "    if is_best:\n",
    "        bestname = os.path.join(save_path, 'model_best.pth.tar')\n",
    "        shutil.copyfile(filename, bestname)\n",
    "\n",
    "\n",
    "def adjust_learning_rate(optimizer, epoch, lr_init):\n",
    "    lr = lr_init * (0 + (1 - 0) * (1 + math.cos(float(epoch) / args.epochs_ft * math.pi)) * 1/2)\n",
    "    for param_group in optimizer.param_groups:\n",
    "        param_group['lr'] = lr\n",
    "    return lr\n",
    "\n",
    "def accuracy(output, target, topk=(1,)):\n",
    "    \"\"\"Computes the precision@k for the specified values of k\"\"\"\n",
    "    maxk = max(topk)\n",
    "    batch_size = target.size(0)\n",
    "\n",
    "    _, pred = output.topk(maxk, 1, True, True)\n",
    "    pred = pred.t()\n",
    "    correct = pred.eq(target.view(1, -1).expand_as(pred))\n",
    "\n",
    "    res = []\n",
    "    for k in topk:\n",
    "        correct_k = correct[:k].view(-1).float().sum(0)\n",
    "        res.append(correct_k.mul_(100.0 / batch_size))\n",
    "    return res\n",
    "\n",
    "\n",
    "def save_obj(obj, name):\n",
    "    with open('obj/' + name + '.pkl', 'wb') as f:\n",
    "        pickle.dump(obj, f, pickle.HIGHEST_PROTOCOL)\n",
    "\n",
    "\n",
    "def load_obj(name):\n",
    "    with open('obj/' + name + '.pkl', 'rb') as f:\n",
    "        return pickle.load(f)\n",
    "\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "\n",
    "if not os.path.isdir(args.save_path):\n",
    "    os.makedirs(args.save_path)\n",
    "log =None\n",
    "\n",
    "# Init dataset\n",
    "if not os.path.isdir(args.data_path):\n",
    "    os.makedirs(args.data_path)\n",
    "\n",
    "if args.dataset == 'cifar10':\n",
    "    mean = [x / 255 for x in [125.3, 123.0, 113.9]]\n",
    "    std = [x / 255 for x in [63.0, 62.1, 66.7]]\n",
    "elif args.dataset == 'cifar100':\n",
    "    mean = [x / 255 for x in [129.3, 124.1, 112.4]]\n",
    "    std = [x / 255 for x in [68.2, 65.4, 70.4]]\n",
    "elif args.dataset=='imagenet':\n",
    "    normalize_imgnet = transforms.Normalize(mean=[0.485, 0.456, 0.406],std=[0.229, 0.224, 0.225])\n",
    "    \n",
    "else:\n",
    "    assert False, \"Unknow dataset : {}\".format(args.dataset)\n",
    "if args.dataset=='imagenet':\n",
    "    pass\n",
    "else:\n",
    "    train_transform = transforms.Compose(\n",
    "        [transforms.RandomHorizontalFlip(), transforms.RandomCrop(32, padding=4), transforms.ToTensor(),\n",
    "         transforms.Normalize(mean, std)])\n",
    "    test_transform = transforms.Compose(\n",
    "        [transforms.ToTensor(), transforms.Normalize(mean, std)])\n",
    "\n",
    "if args.dataset == 'cifar10':\n",
    "    train_data = dset.CIFAR10(args.data_path, train=True, transform=train_transform, download=True)\n",
    "    train_data_test = dset.CIFAR10(args.data_path, train=True, transform=test_transform, download=True)\n",
    "    \n",
    "    test_data = dset.CIFAR10(args.data_path, train=False, transform=test_transform, download=True)\n",
    "    num_classes = 10\n",
    "elif args.dataset == 'imagenet':\n",
    "    \n",
    "    train_data = dset.ImageFolder(os.path.join(args.data_path, 'train'),transforms.Compose([\n",
    "        transforms.RandomSizedCrop(224),\n",
    "        transforms.RandomHorizontalFlip(),\n",
    "        transforms.ToTensor(),\n",
    "        normalize_imgnet,\n",
    "    ]))\n",
    "    train_data_test = dset.ImageFolder(os.path.join(args.data_path, 'train'),transforms.Compose([\n",
    "        transforms.Scale(256), \n",
    "        transforms.CenterCrop(224),\n",
    "        transforms.ToTensor(),\n",
    "        normalize_imgnet,\n",
    "    ]))\n",
    "    \n",
    "    test_data=dset.ImageFolder(os.path.join(args.data_path, 'val'), transforms.Compose([\n",
    "        transforms.Scale(256), \n",
    "        transforms.CenterCrop(224),\n",
    "        transforms.ToTensor(),\n",
    "        normalize_imgnet,\n",
    "    ]))\n",
    "        \n",
    "else:\n",
    "    assert False, 'Do not support dataset : {}'.format(args.dataset)\n",
    "\n",
    "    \n",
    "    \n",
    "#args.batch_size=12    \n",
    "train_loader = torch.utils.data.DataLoader(train_data, batch_size=args.batch_size, shuffle=True,\n",
    "                                           num_workers=args.workers, pin_memory=True)\n",
    "train_loader_kf = torch.utils.data.DataLoader(train_data, batch_size=args.batch_size_kf, shuffle=True,\n",
    "                                           num_workers=args.workers, pin_memory=True)\n",
    "test_loader = torch.utils.data.DataLoader(test_data, batch_size=args.batch_size, shuffle=False,\n",
    "                                          num_workers=args.workers, pin_memory=True)\n",
    "\n",
    "\n",
    "\n",
    "\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "if args.dataset=='cifar10':\n",
    "    netG=generator_cifar.Generator(dim=64)\n",
    "    netG=torch.nn.DataParallel(netG, device_ids=list(range(args.ngpu))).cuda() \n",
    "    netG.load_state_dict(torch.load(os.path.join(args.pretrain_path,'netG_cifar.pth'))) \n",
    "elif args.dataset=='imagenet':\n",
    "    netG=generator_imagenet.ResNetGenerator(64, 128, 4,activation=F.relu, num_classes=0, distribution='normal')\n",
    "    netG=torch.nn.DataParallel(netG, device_ids=list(range(args.ngpu))).cuda()  \n",
    "    netG.module.load_state_dict(torch.load(os.path.join(args.pretrain_path,'netG_imagenet.pth.tar'))['model'])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def train_with_kf(train_loader, model, criterion, optimizer, epoch, log,kfclass):\n",
    "    batch_time = AverageMeter()\n",
    "    data_time = AverageMeter()\n",
    "    losses = AverageMeter()\n",
    "    top1 = AverageMeter()\n",
    "    top5 = AverageMeter()\n",
    "    # switch to train mode\n",
    "    model.train()\n",
    "    \n",
    "    \n",
    "    end = time.time()\n",
    "    for i, (input, target) in enumerate(train_loader):\n",
    "        # measure data loading time\n",
    "        \n",
    "        if input.shape[0]%(args.ngpu)!=0: \n",
    "            continue\n",
    "        \n",
    "        if args.use_cuda:\n",
    "            target = target.cuda()\n",
    "            input = input.cuda()\n",
    "        \n",
    "        data_time.update(time.time() - end)\n",
    "        \n",
    "        if args.dataset=='cifar10':\n",
    "            with torch.no_grad():\n",
    "                kf_input=kfclass(torch.randn(input.shape[0], 128).cuda())\n",
    "        elif args.dataset=='imagenet':\n",
    "            with torch.no_grad():\n",
    "                kf_input=kfclass(torch.empty(input.shape[0], 128, dtype=torch.float32).normal_().cuda())\n",
    "                kf_input=F.interpolate(kf_input,size=224)\n",
    "        \n",
    "        input_list=[]\n",
    "        num_pgpu=input.shape[0]//args.ngpu\n",
    "        for igpu in range(args.ngpu):\n",
    "            input_list.append(torch.cat([input[igpu*num_pgpu:(igpu+1)*num_pgpu],kf_input[igpu*num_pgpu:(igpu+1)*num_pgpu]],dim=0))\n",
    "        input=torch.cat(input_list,dim=0)        \n",
    "\n",
    "\n",
    "        input_var = torch.autograd.Variable(input)\n",
    "        target_var = torch.autograd.Variable(target)\n",
    "\n",
    "\n",
    "        \n",
    "        output = model(input_var)\n",
    "        \n",
    "        output_list=[]\n",
    "        for igpu in range(args.ngpu):\n",
    "            output_list.append(output[igpu*num_pgpu*2:igpu*num_pgpu*2+num_pgpu])\n",
    "        output=torch.cat(output_list,dim=0)  \n",
    "        \n",
    "        loss = criterion(output, target_var)\n",
    "\n",
    "        # measure accuracy and record loss\n",
    "        prec1, prec5 = accuracy(output.data, target, topk=(1, 5))\n",
    "        losses.update(loss.data.item(), input.size(0))\n",
    "        top1.update(prec1.item(), input.size(0))\n",
    "        top5.update(prec5.item(), input.size(0))\n",
    "\n",
    "        # compute gradient and do SGD step\n",
    "        optimizer.zero_grad()\n",
    "        loss.backward()\n",
    "        optimizer.step()\n",
    "        \n",
    "        for module in net.modules():\n",
    "            if isinstance(module,Kf_Conv2d):\n",
    "                module.kfscale.data.clamp_(min=0,max=1) \n",
    "        \n",
    "        \n",
    "        # measure elapsed time\n",
    "        batch_time.update(time.time() - end)\n",
    "        end = time.time()\n",
    "\n",
    "        if i % args.print_freq == 0:\n",
    "            print('  Epoch: [{:03d}][{:03d}/{:03d}]   '\n",
    "                      'Time {batch_time.val:.3f} ({batch_time.avg:.3f})   '\n",
    "                      'Data {data_time.val:.3f} ({data_time.avg:.3f})   '\n",
    "                      'Loss {loss.val:.4f} ({loss.avg:.4f})   '\n",
    "                      'Prec@1 {top1.val:.3f} ({top1.avg:.3f})   '\n",
    "                      'Prec@5 {top5.val:.3f} ({top5.avg:.3f})   '.format(\n",
    "                epoch, i, len(train_loader), batch_time=batch_time,\n",
    "                data_time=data_time, loss=losses, top1=top1, top5=top5) + time_string())\n",
    "    print(\n",
    "        '  **Train** Prec@1 {top1.avg:.3f} Prec@5 {top5.avg:.3f} Error@1 {error1:.3f}'.format(top1=top1, top5=top5,\n",
    "                                                                                              error1=100 - top1.avg),\n",
    "        log)\n",
    "    return top1.avg, losses.avg\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "print(\"=> creating model '{}'\".format(args.arch))\n",
    "\n",
    "\n",
    "if args.dataset=='imagenet':    \n",
    "   \n",
    "    if args.arch=='resnet101':\n",
    "        net=resnet_imagenet.resnet101()\n",
    "    elif args.arch=='resnet50':\n",
    "        net=resnet_imagenet.resnet50()\n",
    "    elif args.arch=='resnet34':\n",
    "        net=resnet_imagenet.resnet34()\n",
    "    elif args.arch=='resnet18':\n",
    "        net=resnet_imagenet.resnet18()\n",
    "else:\n",
    "    if args.arch=='resnet110':\n",
    "        net=models.resnet110(num_classes=10)\n",
    "    elif args.arch=='resnet56':\n",
    "        net=models.resnet56(num_classes=10)\n",
    "    elif args.arch=='resnet32':\n",
    "        net=models.resnet32(num_classes=10)\n",
    "    elif args.arch=='resnet20':\n",
    "        net=models.resnet20(num_classes=10)\n",
    "        \n",
    "\n",
    "if args.dataset=='imagenet':\n",
    "  \n",
    "    if args.arch=='resnet101':\n",
    "        state_dict = torch.load(os.path.join(args.pretrain_path,'resnet101-5d3b4d8f.pth'))       \n",
    "    elif args.arch=='resnet50':\n",
    "        state_dict = torch.load(os.path.join(args.pretrain_path,'resnet50-19c8e357.pth'))\n",
    "    elif args.arch=='resnet34':\n",
    "        state_dict = torch.load(os.path.join(args.pretrain_path,'resnet34-333f7ec4.pth'))\n",
    "    elif args.arch=='resnet18':\n",
    "        state_dict = torch.load(os.path.join(args.pretrain_path,'resnet18-5c106cde.pth'))\n",
    "\n",
    "    from collections import OrderedDict\n",
    "    new_state_dict = OrderedDict()\n",
    "    for k, v in state_dict.items():\n",
    "        if k=='conv1.weight':\n",
    "            new_state_dict['conv1_7x7.weight'] = v\n",
    "        else:\n",
    "            new_state_dict[k] = v    \n",
    "    net.load_state_dict(new_state_dict)         \n",
    "\n",
    "net = torch.nn.DataParallel(net, device_ids=list(range(args.ngpu)))\n",
    "\n",
    "if args.dataset!='imagenet':\n",
    "    if args.arch=='resnet110':\n",
    "        pretrain = torch.load(os.path.join(args.pretrain_path,'cifar_pretrained_nets/','resnet110.pth.tar'))\n",
    "    elif args.arch=='resnet56':\n",
    "        pretrain = torch.load(os.path.join(args.pretrain_path,'cifar_pretrained_nets/','resnet56.pth.tar'))\n",
    "    elif args.arch=='resnet32':\n",
    "        pretrain = torch.load(os.path.join(args.pretrain_path,'cifar_pretrained_nets/','resnet32.pth.tar'))\n",
    "    elif args.arch=='resnet20':\n",
    "        pretrain = torch.load(os.path.join(args.pretrain_path,'cifar_pretrained_nets/','resnet20.pth.tar'))\n",
    "    net.load_state_dict(pretrain['state_dict'].state_dict())\n",
    "\n",
    "\n",
    "\n",
    "\n",
    "\n",
    "\n",
    "\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "net=net.cpu()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "if args.dataset=='imagenet':\n",
    "    def transform_conv( net):\n",
    "  \n",
    "        def _inject(modules):\n",
    "            keys = list(modules.keys())\n",
    "            #print(keys)\n",
    "            for ik, k in enumerate(keys):\n",
    "                if isinstance(modules[k], nn.Conv2d): \n",
    " \n",
    "                    if k!='0' and k!='conv1_7x7': \n",
    "                      \n",
    "                        modules[k] = Kf_Conv2d(modules[k],modules[keys[ik+1]])\n",
    "                        modules[keys[ik+1]]=nn.Sequential() \n",
    "                elif (not isinstance(modules[k], Kf_Conv2d)) and len(modules[k]._modules) > 0: \n",
    "                    _inject(modules[k]._modules)\n",
    "        _inject(net._modules)\n",
    "else:    \n",
    "    def transform_conv( net):\n",
    "\n",
    "        def _inject(modules):\n",
    "            keys = list(modules.keys())\n",
    "          \n",
    "            for ik, k in enumerate(keys):\n",
    "                if isinstance(modules[k], nn.Conv2d): \n",
    "                    if k!='0' and k!='conv_1_3x3': \n",
    "                       \n",
    "                        modules[k] = Kf_Conv2d(modules[k],modules[keys[ik+1]])\n",
    "                        modules[keys[ik+1]]=nn.Sequential() \n",
    "                elif (not isinstance(modules[k], Kf_Conv2d)) and len(modules[k]._modules) > 0: \n",
    "                    _inject(modules[k]._modules)\n",
    "        _inject(net._modules)\n",
    "transform_conv(net) "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "kfconv_list=[]\n",
    "for module in net.modules():\n",
    "    if isinstance(module,Kf_Conv2d):\n",
    "        kfconv_list.append(module)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "kfscale_list=[[] for _ in range(len(kfconv_list))]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "net=net.cuda()\n",
    "\n",
    "criterion = torch.nn.CrossEntropyLoss().cuda()\n",
    "recorder = RecorderMeter(args.epochs)\n",
    "\n",
    "\n",
    "for param in net.parameters(): \n",
    "    param.requires_grad=False\n",
    "for module in net.modules():\n",
    "    if isinstance(module,Kf_Conv2d):\n",
    "        module.kfscale.requires_grad=True\n",
    "    \n",
    "        \n",
    "optimizer = torch.optim.Adam(filter(lambda p: p.requires_grad, net.parameters()), lr=args.learning_rate, betas=(0.5, 0.9))\n",
    "\n",
    "\n",
    "start_time = time.time()\n",
    "epoch_time = AverageMeter()\n",
    "\n",
    "\n",
    "netG.eval()\n",
    "for epoch in range(0, args.epochs):\n",
    "    current_learning_rate =args.learning_rate\n",
    "    need_hour, need_mins, need_secs = convert_secs2time(epoch_time.avg * (args.epochs - epoch))\n",
    "    need_time = '[Need: {:02d}:{:02d}:{:02d}]'.format(need_hour, need_mins, need_secs)\n",
    "\n",
    "    print(\n",
    "        '\\n==>>{:s} [Epoch={:03d}/{:03d}] {:s} [learning_rate={:6.4f}]'.format(time_string(), epoch, args.epochs,\n",
    "        need_time, current_learning_rate) + ' [Best : Accuracy={:.2f}, Error={:.2f}]'.format(recorder.max_accuracy(False),\n",
    "                                                           100 - recorder.max_accuracy(False)))\n",
    "\n",
    "    train_acc, train_los =train_with_kf(train_loader_kf, net, criterion, optimizer, epoch, log,kfclass=netG)\n",
    "\n",
    "    \n",
    "    for ikf in range(len(kfconv_list)):\n",
    "        kfscale_list[ikf].append(kfconv_list[ikf].kfscale.data.clone().cpu())\n",
    "\n",
    "\n",
    "    epoch_time.update(time.time() - start_time)\n",
    "    start_time = time.time()\n",
    "    \n",
    "for param in net.parameters():\n",
    "    param.requires_grad=True "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "for kfscale in kfscale_list[10]:\n",
    "    print(kfscale.squeeze().numpy())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "for kfscale_last in kfscale_list: \n",
    "    print(kfscale_last[-1].squeeze().numpy())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "net=net.cpu()\n",
    "for imd, (nam,module) in enumerate(net.named_modules()):\n",
    "    if isinstance(module, Kf_Conv2d):\n",
    "        module.score=module.bn.weight.data.abs()*(module.kfscale.data-(1-module.kfscale.data)).squeeze() \n",
    "      \n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "for kfconv in kfconv_list:\n",
    "    kfconv.prune_rate=args.prune_rate\n",
    "for imd, (nam,module) in enumerate(net.named_modules()):\n",
    "    if isinstance(module, Kf_Conv2d):\n",
    "        _,index=module.score.sort()\n",
    "        num_pruned_channel=int(module.prune_rate*module.score.shape[0])\n",
    "        \n",
    "        module.out_index=index[num_pruned_channel:]\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "\n",
    "def recover_conv(net):\n",
    "\n",
    "    def _inject(modules):\n",
    "        keys = list(modules.keys())\n",
    "\n",
    "        for ik, k in enumerate(keys):\n",
    "            if isinstance(modules[k], Kf_Conv2d): #### Kf_Conv2d里面没有k=0的\n",
    "                modules[k] =Masked_Conv2d_bn(modules[k])\n",
    "                    \n",
    "            elif (not isinstance(modules[k], Kf_Conv2d)) and len(modules[k]._modules) > 0: # nn.Conv2d的_modules的长度为0，但是Biased_Conv2d的长度为1\n",
    "                _inject(modules[k]._modules)\n",
    "\n",
    "    _inject(net._modules)\n",
    "recover_conv(net)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "net=net.cuda()\n",
    "for input,target in train_loader:\n",
    "    net.train()\n",
    "    with torch.no_grad():\n",
    "        net(input)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "acc_before_ft,acc5_before_ft, loss_before_ft = validate(test_loader, net, criterion, log)\n",
    "\n",
    "\n",
    "optimizer_ft = torch.optim.SGD(net.parameters(), args.lr_ft, momentum=args.momentum,\n",
    "                            weight_decay=args.decay, nesterov=True)\n",
    "\n",
    "recorder_ft = RecorderMeter(args.epochs_ft)\n",
    "\n",
    "recorder_ft_top5=RecorderMeter(args.epochs_ft)\n",
    "\n",
    "\n",
    "for epoch in range(0, args.epochs_ft):\n",
    "    current_learning_rate = adjust_learning_rate(optimizer_ft, epoch, lr_init=args.lr_ft)\n",
    "\n",
    "    need_hour, need_mins, need_secs = convert_secs2time(epoch_time.avg * (args.epochs_ft - epoch))\n",
    "    need_time = '[Need: {:02d}:{:02d}:{:02d}]'.format(need_hour, need_mins, need_secs)\n",
    "\n",
    "    print(\n",
    "        '\\n==>>{:s} [Epoch={:03d}/{:03d}] {:s} [learning_rate={:6.4f}]'.format(time_string(), epoch, args.epochs_ft,\n",
    "                                                                               need_time, current_learning_rate) \\\n",
    "        + ' [Best : Accuracy={:.2f}, Error={:.2f}]'.format(recorder_ft.max_accuracy(False),\n",
    "                                                           100 - recorder_ft.max_accuracy(False)))\n",
    "\n",
    " \n",
    "    train_acc, train_los = train(train_loader, net, criterion, optimizer_ft, epoch, log)\n",
    "\n",
    "   \n",
    "    val_acc_1,val_acc_5, val_los_1 = validate(test_loader, net, criterion, log)\n",
    "\n",
    "\n",
    "    is_best_ft = recorder_ft.update(epoch, train_los, train_acc, val_los_1, val_acc_1)\n",
    "\n",
    "    recorder_ft_top5.update(epoch, train_los, train_acc, val_los_1, val_acc_5)\n",
    "\n",
    "    save_checkpoint({\n",
    "        'epoch': epoch + 1,\n",
    "        'arch': args.arch,\n",
    "        'state_dict': net,\n",
    "        'recorder_ft': recorder_ft,\n",
    "        'optimizer': optimizer_ft.state_dict(),\n",
    "    }, is_best_ft, args.save_path, 'checkpoint_ft.pth.tar')\n",
    "\n",
    "    # measure elapsed time\n",
    "    epoch_time.update(time.time() - start_time)\n",
    "    start_time = time.time() "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "print('top1:',recorder_ft.max_accuracy(False))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "masked_conv_list=[]\n",
    "for imd, (nam,module) in enumerate(net.named_modules()):\n",
    "    if isinstance(module, Masked_Conv2d_bn):\n",
    "        masked_conv_list.append((nam,module))\n",
    "if args.dataset=='imagenet':\n",
    "    for imd in range(len(masked_conv_list)):\n",
    "        \n",
    "        if 'conv2' in masked_conv_list[imd][0] or 'conv3' in masked_conv_list[imd][0]:\n",
    "            masked_conv_list[imd][1].in_index=masked_conv_list[imd-1][1].out_index\n",
    "else:\n",
    "    for imd in range(len(masked_conv_list)):\n",
    "        \n",
    "        if 'conv_b' in masked_conv_list[imd][0]:\n",
    "            masked_conv_list[imd][1].in_index=masked_conv_list[imd-1][1].out_index"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "\n",
    "\n",
    "if args.dataset=='imagenet':\n",
    "    def pruning_conv( net):\n",
    "       \n",
    "        def _inject(modules):\n",
    "            keys = list(modules.keys())\n",
    "       \n",
    "            for ik, k in enumerate(keys):\n",
    "                if isinstance(modules[k], Masked_Conv2d_bn): \n",
    "                 \n",
    "                    if args.arch=='resnet18' or args.arch=='resnet34':\n",
    "                        if 'conv1' in k:\n",
    "                            modules[k] = Pruned_Conv2d_bn1(modules[k])\n",
    "                        elif 'conv2' in k:      \n",
    "                            modules[k] = Pruned_Conv2d_bn2(modules[k])\n",
    "                    else: ##### bottleneck结构\n",
    "                        if 'conv1' in k:\n",
    "                            modules[k] = Pruned_Conv2d_bn1(modules[k])\n",
    "                        elif 'conv2' in k:      \n",
    "                            modules[k] = Pruned_Conv2d_bn_middle(modules[k])\n",
    "                        elif 'conv3' in k:      \n",
    "                            modules[k] = Pruned_Conv2d_bn2(modules[k])\n",
    "\n",
    "                elif (not isinstance(modules[k], Kf_Conv2d)) and len(modules[k]._modules) > 0: \n",
    "                    _inject(modules[k]._modules)\n",
    "        _inject(net._modules)\n",
    "\n",
    "else:\n",
    "    def pruning_conv( net):\n",
    "     \n",
    "        def _inject(modules):\n",
    "            keys = list(modules.keys())\n",
    "       \n",
    "            for ik, k in enumerate(keys):\n",
    "                if isinstance(modules[k], Masked_Conv2d_bn): \n",
    "                  \n",
    "                    if 'conv_a' in k:\n",
    "                        modules[k] = Pruned_Conv2d_bn1(modules[k])\n",
    "                    elif 'conv_b' in k:      \n",
    "                        modules[k] = Pruned_Conv2d_bn2(modules[k])\n",
    "\n",
    "                elif (not isinstance(modules[k], Kf_Conv2d)) and len(modules[k]._modules) > 0: \n",
    "                    _inject(modules[k]._modules)\n",
    "        _inject(net._modules)\n",
    "pruning_conv(net)"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.6.5"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
